import akshare as ak
import pandas as pd
import numpy as np
import talib
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, MinMaxScaler
import requests
from bs4 import BeautifulSoup
import re
import yfinance as yf

# 获取股票价格和成交量数据
def get_stock_data(symbol):
    try:
        stock_data = ak.stock_zh_a_daily(symbol=symbol)
        return stock_data
    except Exception as e:
        print(f"Error fetching data: {e}")
        return pd.DataFrame()

# 计算技术指标
def calculate_technical_indicators(data):
    data['ma5'] = talib.SMA(data['close'], timeperiod=5)
    data['ma10'] = talib.SMA(data['close'], timeperiod=10)
    data['ma20'] = talib.SMA(data['close'], timeperiod=20)
    data['rsi'] = talib.RSI(data['close'], timeperiod=14)
    data['macd'], data['macd_signal'], data['macd_hist'] = talib.MACD(data['close'], fastperiod=12, slowperiod=26, signalperiod=9)
    data['upper_band'], data['middle_band'], data['lower_band'] = talib.BBANDS(data['close'], timeperiod=20)
    return data

# 获取基本面数据
def get_fundamental_data(symbol):
    stock = yf.Ticker(symbol)
    info = stock.info
    pe_ratio = info.get('trailingPE', None)
    pb_ratio = info.get('priceToBook', None)
    return pe_ratio, pb_ratio

# 模拟获取市场情绪数据（这里通过简单的新闻标题关键词计数来模拟）
def get_market_sentiment(symbol):
    url = f"https://so.toutiao.com/search?dvpf=pc&source=sug&keyword={symbol}&page_num=0&pd=synthesis"
    response = requests.get(url)
    soup = BeautifulSoup(response.text, 'html.parser')
    news_titles = [title.text for title in soup.find_all('h2', class_='news-title')]
    positive_keywords = ['上涨', '利好', '突破']
    negative_keywords = ['下跌', '亏损', '危机']
    positive_count = sum([len(re.findall(keyword, title)) for keyword in positive_keywords for title in news_titles])
    negative_count = sum([len(re.findall(keyword, title)) for keyword in negative_keywords for title in news_titles])
    return positive_count - negative_count

def prepare_data(stock_data, symbol):
    # Check if the symbol exists in the stock_data columns
    if symbol not in stock_data.columns:
        raise KeyError(f"Symbol {symbol} not found in stock_data columns")

    # Assuming technical_data is derived from stock_data
    technical_data = stock_data[symbol]

    # Ensure all arrays have the same number of rows
    required_columns = ['ma5', 'ma10', 'ma20', 'rsi', 'macd', 'macd_signal', 'macd_hist', 'upper_band', 'lower_band']
    min_length = min(len(technical_data[col]) for col in required_columns)

    # Truncate all columns to the minimum length
    truncated_data = {col: technical_data[col][:min_length] for col in required_columns}

    # Stack the truncated columns
    X = np.column_stack([truncated_data[col] for col in required_columns])

    # Assuming y is derived similarly
    y = stock_data['target'][:min_length]

    return X, y

# 训练模型
def train_model(X, y):
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train_scaled, y_train)
    return model

# 预测涨停概率
def predict_limit_up(model, X):
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    probabilities = model.predict_proba(X_scaled)[:, 1]
    return probabilities

def main():
    symbol = "sh600000"  # 这里假设一个股票代码，实际可根据需求调整
    stock_data = get_stock_data(symbol)
    if not stock_data.empty:
        X, y = prepare_data(stock_data, symbol)
        model = train_model(X, y)
        latest_data = stock_data.tail(1)
        X_latest, _ = prepare_data(latest_data, symbol)
        probability = predict_limit_up(model, X_latest)
        print(f"涨停概率: {probability[0]}")
    else:
        print("No data available.")

if __name__ == '__main__':
    main()